import argparse, time, toml

from threading import Thread
from .collect import *
from .collect_bcc_perf import *
from .collect_node import *
from .collect_ctr import *


def prepare_output_dir(output_dir):
    if os.path.exists(output_dir):
        os.system(
            " ".join(
                [
                    "mv",
                    output_dir,
                    output_dir + time.strftime("%Y_%m_%d_%H_%M_%S", time.localtime()),
                ]
            )
        )
    os.system(
        " ".join(
            [
                "mkdir -p",
                output_dir,
            ]
        )
    )


def setup_bcc(args, output):
    collect = BccCollect(
        args["path"],
        args["probes"],
        args["timeout"],
        output,
        args["interval"],
    )
    collect.collect()


def setup_perfstat(args, output):
    collect = PerfStatCollect(args["timeout"], args["interval"], output, args["events"])
    collect.collect()


# TODO: reset interval?
def setup_ctr(args, output):
    """collect info for each container"""
    collect = CtrMainCollect(
        args["interval"],
        output,
        args["collects"],
        args["driver"],
        args["cgroupPrefix"],
    )
    collect.collect()


def setup_node(args, output):
    collect = NodeCollect(
        output, args["interval"], args["collects"], args["cgroupPrefix"], args["pid"]
    )
    collect.collect()


def setup(configs):
    prepare_output_dir(configs["output"])

    # print("configs", configs)
    # TODO: set default value for each collect
    threads = []
    if configs["bcc"]:
        bcc_thread = Thread(target=setup_bcc, args=(configs["bcc"], configs["output"]))
        bcc_thread.start()
        threads.append(bcc_thread)

    if configs["perfstat"]:
        perfstat_thread = Thread(
            target=setup_perfstat, args=(configs["perfstat"], configs["output"])
        )
        perfstat_thread.start()
        threads.append(perfstat_thread)

    if configs["ctr"]:
        ctr_thread = Thread(target=setup_ctr, args=(configs["ctr"], configs["output"]))
        ctr_thread.start()
        threads.append(ctr_thread)

    if configs["node"]:
        node_thread = Thread(
            target=setup_node, args=(configs["node"], configs["output"])
        )
        node_thread.start()
        threads.append(node_thread)

    # TODO: wait timeout and exit?
    # TODO: stress-ng
    for thread in threads:
        thread.join()


def main() -> None:
    """
    usage example: `python -m collect  -f /path/to/config.toml  -t time`
    -t: total time to run. No timeout by default
    """
    parser = argparse.ArgumentParser(
        description="A tool for metrics collection of hardware and software workloads",
        epilog=(
            "Through the software-hardware collaborative analysis method, the application"
            " characteristics and system-level characteristics of load execution under the"
            " conditions of differnet configurations of clusters and different co-location"
            " modes are studied to guide the resource planning and scheduling management of"
            " cloud clusters"
        ),
        add_help=True,
    )

    parser.add_argument(
        "-f",
        "--config",
        type=str,
        help="config file path",
        default="./collect_py/config.toml",
        required=False,
    )

    parser.add_argument(
        "-t",
        "--timeout",
        type=str,
        help="time to run collect",
        default="",  # no timeout by default
        required=False,
    )
    args = parser.parse_args()
    print("args", args)
    setup(toml.load(args.config))


if __name__ == "__main__":
    main()
